﻿/* ===================================尊重原著======================================
 * 创 建 者：mao xiao jun 
 * 创建日期：2021-12-12
 * CLR Version : any
 * ===================================尊重原著=====================================*/
using System;
using System.Collections.Concurrent;
using System.IO;
using System.Net.Sockets;
using System.Threading;
using System.Threading.Tasks;

namespace MTP
{
    class MtpClient : MtpDispatch, IMtpClient
    {

        private readonly TcpClient client;
        private readonly string hostname;
        private readonly int port;

        private readonly ConcurrentDictionary<Guid, NotifyLock<Stream>> notifyHandler = new ConcurrentDictionary<Guid, NotifyLock<Stream>>();

        public MtpClient(string hostname, int port, string cipher) : base(cipher)
        {
            client = new TcpClient();
            this.hostname = hostname;
            this.port = port;
        }

        public bool IsConnected => client.Connected;


        public void Connect()
        {
            ConnectAsync().GetAwaiter().GetResult();
        }

        public async Task ConnectAsync()
        {
            await client.ConnectAsync(hostname, port);
            var stream = client.GetStream();
            await SendHandshake(stream);
            if (!await VerifyHandshake(stream))
            {
                throw new InvalidDataException("Invalid secret key");
            }
            Receive().GetAwaiter();
        }

        public Task<Stream> SendAsync(Stream stream, string header)
        {
            return Task.Run(() =>
            {
                Guid guid = Guid.NewGuid();
                try
                {
                    NotifyLock<Stream> notifyLock = new NotifyLock<Stream>();
                    notifyHandler.TryAdd(guid, notifyLock);
                    SendData(guid, client.GetStream(), stream, header).GetAwaiter().GetResult();
                    if (!notifyLock.Wait(60000))
                        throw new TimeoutException("MTP 响应超时");
                    return notifyLock.Value;
                }
                finally
                {
                    notifyHandler.TryRemove(guid, out _);
                }
            });
        }

        public async Task<byte[]> SendAsync(byte[] data, string header)
        {
            using var ms = new MemoryStream(data);
            using var stream = await SendAsync(ms, header);
            using var ms2 = new MemoryStream();
            await stream.CopyToAsync(ms2);
            return ms2.ToArray();
        }


        private async Task Receive()
        {
            while (client.Connected)
            {
                NetworkStream stream = client.GetStream();
                try
                {
                    await ReadNext(stream, default);
                }
                catch (Exception)
                {
                    client.Close();
                    break;
                }
            }
            OnDisconnectEvent?.Invoke();
            client.Dispose();
        }

        public IMtpClient OnDisconnect(Action onDisconnectEvent)
        {
            OnDisconnectEvent += onDisconnectEvent;
            return this;
        }

        public event Action OnDisconnectEvent;

        public void Close()
        {
            client.Close();
            client.Dispose();
        }

        protected override Task ReadData(Guid guid, Stream stream, string header, NetworkStream networkStream)
        {
            if (notifyHandler.TryGetValue(guid, out var notifyLock))
            {
                notifyLock.Notify(stream);
            }
            return Task.CompletedTask;
        }

        class NotifyLock<T>
        {
            private readonly ManualResetEvent manualResetEvent;
            public NotifyLock()
            {
                manualResetEvent = new ManualResetEvent(false);
            }

            public bool Wait(int millisecondsTimeout)
            {
                return manualResetEvent.WaitOne(millisecondsTimeout);
            }

            public T Value { get; set; }

            public void Notify(T value)
            {
                Value = value;
                manualResetEvent.Set();
            }

        }
    }
}
